import os
import aiofiles
import mimetypes
from typing import AsyncGenerator, Any, BinaryIO, Optional, Dict, Union
from fastapi import UploadFile
from config import TEMP_DIR

# 扩展MIME类型映射
MIME_TYPE_MAP = {
    # 文本类型
    'text': 'text/plain',
    'json': 'application/json',
    'html': 'text/html',
    'xml': 'application/xml',
    # 音频类型
    'mp3': 'audio/mpeg',
    'wav': 'audio/wav',
    'ogg': 'audio/ogg',
    'flac': 'audio/flac',
    # 视频类型
    'mp4': 'video/mp4',
    'webm': 'video/webm',
    'avi': 'video/x-msvideo',
    'mkv': 'video/x-matroska',
    # 图片类型
    'jpg': 'image/jpeg',
    'jpeg': 'image/jpeg',
    'png': 'image/png',
    'gif': 'image/gif',
    'webp': 'image/webp',
    'svg': 'image/svg+xml'
}

async def save_temp_file(upload_file: UploadFile) -> str:
    """Save an uploaded file to a temporary location."""
    temp_file = os.path.join(TEMP_DIR, f"temp_{upload_file.filename}")
    async with aiofiles.open(temp_file, 'wb') as f:
        await f.write(await upload_file.read())
    return temp_file

def remove_temp_file(temp_file: str):
    """Remove a temporary file."""
    if os.path.exists(temp_file):
        os.remove(temp_file)

async def read_file_stream(file_path: str, chunk_size: int = 1024, media_type: Optional[str] = None) -> AsyncGenerator[bytes, None]:
    """Read a file in chunks as an async generator.
    
    Args:
        file_path: Path to the file to read
        chunk_size: Size of chunks to read
        media_type: Optional media type to use for processing
        
    Returns:
        AsyncGenerator yielding bytes
    """
    if not media_type:
        media_type = get_mime_type(file_path)
    
    async with aiofiles.open(file_path, 'rb') as f:
        while chunk := await f.read(chunk_size):
            # 根据媒体类型进行特定处理
            if media_type.startswith('text/'):
                # 确保文本编码正确
                try:
                    # 尝试解码并重新编码以确保一致性
                    decoded = chunk.decode('utf-8')
                    chunk = decoded.encode('utf-8')
                except UnicodeDecodeError:
                    # 如果解码失败，保持原样
                    pass
            # 对于音频和视频流，保持原始二进制格式
            yield chunk

async def write_file_stream(file_path: str, data_stream: AsyncGenerator[Any, None], 
                           chunk_size: int = 1024, media_type: Optional[str] = None) -> AsyncGenerator[bytes, None]:
    """Write a stream to a file and yield progress.
    
    Args:
        file_path: Path to write the file to
        data_stream: Stream of data to write
        chunk_size: Size of chunks to write
        media_type: Optional media type to use for processing
        
    Returns:
        AsyncGenerator yielding bytes written
    """
    if not media_type:
        media_type = get_mime_type(file_path)
        
    async with aiofiles.open(file_path, 'wb') as f:
        async for chunk in data_stream:
            # 根据数据类型和媒体类型进行处理
            if isinstance(chunk, str):
                if media_type.startswith('text/') or media_type == 'application/json':
                    chunk = chunk.encode('utf-8')
                else:
                    # 非文本类型但收到字符串，尝试解码为二进制
                    try:
                        import base64
                        chunk = base64.b64decode(chunk)
                    except Exception:
                        # 如果解码失败，使用UTF-8编码
                        chunk = chunk.encode('utf-8')
            
            # 确保chunk是bytes类型
            if not isinstance(chunk, bytes):
                chunk = str(chunk).encode('utf-8')
                
            await f.write(chunk)
            yield chunk

def get_file_type(file_path: str) -> Optional[str]:
    """Determine file type based on extension and mime type.
    
    Args:
        file_path: Path to the file
        
    Returns:
        String indicating file type ('image', 'audio', 'video', 'text') or None
    """
    # 先从文件扩展名获取
    _, ext = os.path.splitext(file_path)
    ext = ext.lower().lstrip('.')
    
    # 使用扩展的MIME类型映射
    if ext in MIME_TYPE_MAP:
        mime_type = MIME_TYPE_MAP[ext]
    else:
        # 回退到系统的mimetypes
        mime_type, _ = mimetypes.guess_type(file_path)
    
    if mime_type:
        if mime_type.startswith('image/'):
            return 'image'
        elif mime_type.startswith('audio/'):
            return 'audio'
        elif mime_type.startswith('video/'):
            return 'video'
        elif mime_type.startswith('text/') or mime_type == 'application/json':
            return 'text'
        elif mime_type == 'application/octet-stream':
            # 尝试通过文件扩展名进一步判断
            if ext in ['mp3', 'wav', 'ogg', 'flac']:
                return 'audio'
            elif ext in ['mp4', 'webm', 'avi', 'mkv']:
                return 'video'
            elif ext in ['jpg', 'jpeg', 'png', 'gif', 'webp', 'svg']:
                return 'image'
            elif ext in ['txt', 'json', 'html', 'xml']:
                return 'text'
    return None

def get_mime_type(file_path: str) -> str:
    """Get the MIME type of a file.
    
    Args:
        file_path: Path to the file
        
    Returns:
        MIME type string
    """
    # 先从文件扩展名获取
    _, ext = os.path.splitext(file_path)
    ext = ext.lower().lstrip('.')
    
    # 使用扩展的MIME类型映射
    if ext in MIME_TYPE_MAP:
        return MIME_TYPE_MAP[ext]
    
    # 回退到系统的mimetypes
    mime_type, _ = mimetypes.guess_type(file_path)
    return mime_type or 'application/octet-stream'

async def create_stream_response(data_stream: AsyncGenerator[Any, None], media_type: str, 
                               filename: Optional[str] = None) -> AsyncGenerator[bytes, None]:
    """创建适合流式响应的生成器，根据媒体类型设置正确的处理逻辑
    
    Args:
        data_stream: 数据流
        media_type: 媒体类型
        filename: 可选的文件名，用于下载响应
        
    Returns:
        处理后的字节流
    """
    headers = {}
    if filename:
        content_disposition = f'attachment; filename="{filename}"'
        headers['Content-Disposition'] = content_disposition
    
    async for chunk in data_stream:
        # 确保chunk是bytes类型
        if isinstance(chunk, str):
            if media_type.startswith('text/') or media_type == 'application/json':
                chunk = chunk.encode('utf-8')
            else:
                # 非文本类型但收到字符串，尝试解码为二进制
                try:
                    import base64
                    chunk = base64.b64decode(chunk)
                except Exception:
                    chunk = chunk.encode('utf-8')
        elif not isinstance(chunk, bytes):
            chunk = str(chunk).encode('utf-8')
            
        yield chunk